from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from rag.knowledge_base_manager import KnowledgeBaseManager
from monitoring.decorators import monitor_api_endpoint, monitor_rag_query, monitor_document_upload
import os
import tempfile
import logging
from config.thresholds import RAGThresholds, KnowledgeBaseThresholds, ValidationThresholds

logger = logging.getLogger(__name__)

# 创建路由器
router = APIRouter()

# 初始化知识库管理器
kb_manager = KnowledgeBaseManager()

# 请求和响应模型
class CreateKnowledgeBaseRequest(BaseModel):
    name: str
    description: str = ""
    chunk_size: int = KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE
    chunk_overlap: int = KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP
    splitter_type: str = "recursive"
    custom_separators: Optional[List[str]] = None
    metadata_filters: Optional[Dict[str, Any]] = None
    session_id: Optional[str] = None
    user_id: Optional[str] = None

class KnowledgeBaseInfo(BaseModel):
    name: str
    description: str
    chunk_size: int
    chunk_overlap: int
    splitter_type: str
    document_count: int

class UploadToKBRequest(BaseModel):
    knowledge_base_name: str
    custom_metadata: Optional[Dict[str, Any]] = None

class UploadToKBResponse(BaseModel):
    status: str
    message: str
    knowledge_base: str
    collection_info: Dict[str, Any]
    file_path: str
    chunks_created: int

class QuestionInKBRequest(BaseModel):
    question: str
    knowledge_base_name: str
    top_k: Optional[int] = RAGThresholds.DEFAULT_TOP_K
    rerank_threshold: Optional[float] = RAGThresholds.DEFAULT_RERANK_THRESHOLD
    session_id: Optional[str] = None
    user_id: Optional[str] = None

class QuestionInKBResponse(BaseModel):
    status: str
    question: str
    answer: str
    sources: List[Dict[str, Any]]
    context_length: int
    documents_retrieved: int
    knowledge_base: str
    kb_config: Dict[str, str]

class KnowledgeBaseListResponse(BaseModel):
    status: str
    knowledge_bases: List[KnowledgeBaseInfo]
    count: int

# 创建知识库
@router.post("/create", response_model=Dict[str, Any])
@monitor_api_endpoint("kb_create_knowledge_base")
async def create_knowledge_base(request: CreateKnowledgeBaseRequest):
    """创建新的知识库"""
    try:
        logger.info(f"创建知识库请求: {request.name}")
        
        result = kb_manager.create_knowledge_base(
            name=request.name,
            description=request.description,
            chunk_size=request.chunk_size,
            chunk_overlap=request.chunk_overlap,
            splitter_type=request.splitter_type,
            custom_separators=request.custom_separators,
            metadata_filters=request.metadata_filters
        )
        
        if result["status"] == "success":
            return result
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"创建知识库失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"创建知识库失败: {str(e)}")

# 上传文档到知识库
@router.post("/upload", response_model=UploadToKBResponse)
@monitor_api_endpoint("kb_upload_document")
@monitor_document_upload()
async def upload_document_to_kb(
    file: UploadFile = File(...),
    knowledge_base_name: str = Form(...),
    custom_metadata: Optional[str] = Form(None),
    auto_create: bool = Form(True),
    kb_description: Optional[str] = Form(None),
    chunk_size: Optional[int] = Form(None),
    chunk_overlap: Optional[int] = Form(None),
    splitter_type: Optional[str] = Form("recursive"),
    session_id: Optional[str] = Form(None),
    user_id: Optional[str] = Form(None)
):
    """上传文档到指定知识库"""
    try:
        logger.info(f"上传文档到知识库: {knowledge_base_name}, 文件: {file.filename}, 自动创建: {auto_create}")
        
        # 解析自定义元数据
        metadata = None
        if custom_metadata:
            import json
            try:
                metadata = json.loads(custom_metadata)
            except:
                metadata = {"custom_metadata": custom_metadata}
        
        # 构建知识库配置
        kb_config = None
        if auto_create:
            kb_config = {
                "description": kb_description or f"自动创建的知识库: {knowledge_base_name}",
                "chunk_size": chunk_size or KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE,
                "chunk_overlap": chunk_overlap or KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP,
                "splitter_type": splitter_type or "recursive"
            }
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp_file:
            content = await file.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # 上传到知识库
        result = kb_manager.upload_document_to_kb(
            file_path=temp_file_path,
            knowledge_base_name=knowledge_base_name,
            custom_metadata=metadata,
            auto_create=auto_create,
            kb_config=kb_config
        )
        
        # 删除临时文件
        os.unlink(temp_file_path)
        
        if result["status"] == "success":
            return UploadToKBResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"上传文档失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"上传文档失败: {str(e)}")

# 智能上传文档（自动创建知识库）
@router.post("/smart-upload", response_model=UploadToKBResponse)
@monitor_api_endpoint("kb_smart_upload_document")
@monitor_document_upload()
async def smart_upload_document(
    file: UploadFile = File(...),
    knowledge_base_name: str = Form(...),
    custom_metadata: Optional[str] = Form(None),
    kb_description: Optional[str] = Form(None),
    chunk_size: Optional[int] = Form(None),
    chunk_overlap: Optional[int] = Form(None),
    splitter_type: Optional[str] = Form("recursive"),
    session_id: Optional[str] = Form(None),
    user_id: Optional[str] = Form(None)
):
    """智能上传文档：如果知识库不存在则自动创建"""
    try:
        logger.info(f"智能上传文档到知识库: {knowledge_base_name}, 文件: {file.filename}")
        
        # 解析自定义元数据
        metadata = None
        if custom_metadata:
            import json
            try:
                metadata = json.loads(custom_metadata)
            except:
                metadata = {"custom_metadata": custom_metadata}
        
        # 构建知识库配置
        kb_config = {
            "description": kb_description or f"智能创建的知识库: {knowledge_base_name}",
            "chunk_size": chunk_size or KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE,
            "chunk_overlap": chunk_overlap or KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP,
            "splitter_type": splitter_type or "recursive"
        }
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp_file:
            content = await file.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # 智能上传到知识库
        result = kb_manager.upload_document_to_kb(
            file_path=temp_file_path,
            knowledge_base_name=knowledge_base_name,
            custom_metadata=metadata,
            auto_create=True,  # 强制自动创建
            kb_config=kb_config
        )
        
        # 删除临时文件
        os.unlink(temp_file_path)
        
        if result["status"] == "success":
            return UploadToKBResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"智能上传文档失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"智能上传文档失败: {str(e)}")

# 在知识库中回答问题
@router.post("/ask", response_model=QuestionInKBResponse)
@monitor_api_endpoint("kb_ask_question")
@monitor_rag_query()
async def ask_question_in_kb(request: QuestionInKBRequest):
    """在指定知识库中回答问题"""
    try:
        logger.info(f"在知识库 {request.knowledge_base_name} 中提问: {request.question}")
        
        result = kb_manager.ask_question_in_kb(
            question=request.question,
            knowledge_base_name=request.knowledge_base_name,
            top_k=request.top_k,
            rerank_threshold=request.rerank_threshold,
            user_id=request.user_id,
            session_id=request.session_id
        )
        
        if result["status"] == "success":
            return QuestionInKBResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"回答问题失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"回答问题失败: {str(e)}")

# 列出所有知识库
@router.get("/list", response_model=KnowledgeBaseListResponse)
@monitor_api_endpoint("kb_list_knowledge_bases")
async def list_knowledge_bases():
    """列出所有知识库"""
    try:
        result = kb_manager.list_knowledge_bases()
        
        if result["status"] == "success":
            return KnowledgeBaseListResponse(**result)
        else:
            raise HTTPException(status_code=400, detail=result["message"])
            
    except Exception as e:
        logger.error(f"列出知识库失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"列出知识库失败: {str(e)}")

# 列出某个知识库中的文件
@router.get("/{knowledge_base_name}/files")
@monitor_api_endpoint("kb_list_files")
async def list_kb_files(knowledge_base_name: str):
    """列出某个知识库中的所有文件"""
    try:
        logger.info(f"列出知识库文件: {knowledge_base_name}")
        
        # 获取知识库信息
        kb_list = kb_manager.list_knowledge_bases()
        if kb_list["status"] != "success":
            return {
                "status": "error",
                "message": kb_list["message"],
                "knowledge_base": knowledge_base_name,
                "files": [],
                "total_files": 0,
                "total_chunks": 0
            }
        
        # 查找知识库
        kb_found = False
        for kb in kb_list["knowledge_bases"]:
            if kb["name"] == knowledge_base_name:
                kb_found = True
                break
        
        if not kb_found:
            return {
                "status": "error",
                "message": f"知识库 '{knowledge_base_name}' 不存在",
                "knowledge_base": knowledge_base_name,
                "files": [],
                "total_files": 0,
                "total_chunks": 0
            }
        
        # 获取集合信息
        collection_info = kb_manager.document_processor.get_collection_info(knowledge_base_name)
        
        # 搜索所有文档来获取文件列表
        try:
            all_docs = kb_manager.document_processor.search_similar(
                "", 
                collection_name=knowledge_base_name, 
                top_k=1000  # 获取所有文档
            )
            
            # 按文件分组
            files = {}
            for doc in all_docs:
                file_name = doc.metadata.get("file_name", "未知文件")
                source = doc.metadata.get("source", "未知来源")
                file_type = doc.metadata.get("file_type", "未知类型")
                
                if file_name not in files:
                    files[file_name] = {
                        "file_name": file_name,
                        "source": source,
                        "file_type": file_type,
                        "chunks_count": 0,
                        "upload_time": doc.metadata.get("upload_time", "未知")
                    }
                files[file_name]["chunks_count"] += 1
            
            return {
                "status": "success",
                "knowledge_base": knowledge_base_name,
                "files": list(files.values()),
                "total_files": len(files),
                "total_chunks": collection_info["document_count"]
            }
            
        except Exception as e:
            logger.warning(f"获取文件列表失败，返回基本信息: {e}")
            return {
                "status": "success",
                "knowledge_base": knowledge_base_name,
                "files": [],
                "total_files": 0,
                "total_chunks": collection_info["document_count"],
                "message": "无法获取详细文件列表，但知识库存在"
            }
            
    except Exception as e:
        logger.error(f"列出知识库文件失败: {str(e)}")
        return {
            "status": "error",
            "message": f"列出知识库文件失败: {str(e)}",
            "knowledge_base": knowledge_base_name,
            "files": [],
            "total_files": 0,
            "total_chunks": 0
        }

# 删除知识库中的特定文件
@router.delete("/{knowledge_base_name}/files/{file_name}")
@monitor_api_endpoint("kb_delete_file")
async def delete_kb_file(knowledge_base_name: str, file_name: str):
    """删除知识库中的特定文件"""
    try:
        logger.info(f"删除知识库文件: {knowledge_base_name}/{file_name}")
        
        # 检查知识库是否存在
        kb_list = kb_manager.list_knowledge_bases()
        if kb_list["status"] != "success":
            return {
                "status": "error",
                "message": kb_list["message"],
                "deleted_chunks": 0,
                "remaining_chunks": 0
            }
        
        kb_found = False
        for kb in kb_list["knowledge_bases"]:
            if kb["name"] == knowledge_base_name:
                kb_found = True
                break
        
        if not kb_found:
            return {
                "status": "error",
                "message": f"知识库 '{knowledge_base_name}' 不存在",
                "deleted_chunks": 0,
                "remaining_chunks": 0
            }
        
        # 获取所有文档
        all_docs = kb_manager.document_processor.search_similar(
            "", 
            collection_name=knowledge_base_name, 
            top_k=1000
        )
        
        # 找到要删除的文件的所有文档块
        docs_to_delete = []
        for doc in all_docs:
            if doc.metadata.get("file_name") == file_name:
                docs_to_delete.append(doc)
        
        if not docs_to_delete:
            return {
                "status": "error",
                "message": f"文件 '{file_name}' 在知识库中不存在",
                "deleted_chunks": 0,
                "remaining_chunks": len(all_docs)
            }
        
        # 从向量存储中删除文档
        # 注意：ChromaDB的删除功能需要文档ID，这里我们使用一个简化的方法
        # 在实际应用中，你可能需要维护一个文档ID到文件名的映射
        
        # 重新创建集合，排除要删除的文件
        remaining_docs = [doc for doc in all_docs if doc.metadata.get("file_name") != file_name]
        
        if remaining_docs:
            # 重新创建集合
            kb_manager.document_processor.create_vector_store(
                documents=remaining_docs,
                collection_name=knowledge_base_name
            )
        else:
            # 如果删除后没有文档了，创建空集合
            kb_manager.document_processor.create_vector_store(
                documents=[],
                collection_name=knowledge_base_name
            )
        
        return {
            "status": "success",
            "message": f"成功删除文件 '{file_name}' 从知识库 '{knowledge_base_name}'",
            "deleted_chunks": len(docs_to_delete),
            "remaining_chunks": len(remaining_docs)
        }
        
    except Exception as e:
        logger.error(f"删除知识库文件失败: {str(e)}")
        return {
            "status": "error",
            "message": f"删除知识库文件失败: {str(e)}",
            "deleted_chunks": 0,
            "remaining_chunks": 0
        }

# 删除知识库
@router.delete("/{knowledge_base_name}")
@monitor_api_endpoint("kb_delete_knowledge_base")
async def delete_knowledge_base(knowledge_base_name: str):
    """删除知识库"""
    try:
        logger.info(f"删除知识库: {knowledge_base_name}")
        
        result = kb_manager.delete_knowledge_base(knowledge_base_name)
        
        # 直接返回结果，不抛出HTTPException
        return result
            
    except Exception as e:
        logger.error(f"删除知识库失败: {str(e)}")
        return {
            "status": "error",
            "message": f"删除知识库失败: {str(e)}"
        }

# 获取知识库信息
@router.get("/{knowledge_base_name}/info")
@monitor_api_endpoint("kb_get_info")
async def get_knowledge_base_info(knowledge_base_name: str):
    """获取知识库详细信息"""
    try:
        kb_list = kb_manager.list_knowledge_bases()
        
        if kb_list["status"] == "success":
            for kb in kb_list["knowledge_bases"]:
                if kb["name"] == knowledge_base_name:
                    return {
                        "status": "success",
                        "knowledge_base": kb
                    }
            
            # 知识库不存在，返回200但status为error
            return {
                "status": "error",
                "message": f"知识库 '{knowledge_base_name}' 不存在",
                "knowledge_base": None
            }
        else:
            return {
                "status": "error",
                "message": kb_list["message"],
                "knowledge_base": None
            }
            
    except Exception as e:
        logger.error(f"获取知识库信息失败: {str(e)}")
        return {
            "status": "error",
            "message": f"获取知识库信息失败: {str(e)}",
            "knowledge_base": None
        }

# 健康检查
@router.get("/health")
@monitor_api_endpoint("kb_health_check")
async def health_check():
    """知识库管理系统健康检查"""
    return {
        "status": "healthy",
        "service": "Knowledge Base Manager",
        "features": [
            "create_knowledge_base",
            "upload_documents", 
            "ask_questions",
            "list_knowledge_bases",
            "delete_knowledge_bases"
        ]
    } 